package com.ml4ai.backend.stack.webcollector;

import cn.edu.hfut.dmic.webcollector.crawldb.DBManager;
import cn.edu.hfut.dmic.webcollector.crawldb.Generator;
import cn.edu.hfut.dmic.webcollector.model.CrawlDatum;
import cn.edu.hfut.dmic.webcollector.model.CrawlDatums;
import com.google.gson.Gson;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * Created by leecheng on 2018/10/19.
 */
public class StorageAdapterDBManager extends DBManager {

    Logger LOG = LoggerFactory.getLogger(StorageAdapterDBManager.class);
    public StoreDB storeDB;
    private Gson gson = new Gson();

    public StorageAdapterDBManager(Map<String, String> crawl, Map<String, String> fetch, Map<String, String> link, Map<String, String> redirect) {
        this.storeDB = new StoreDB(crawl, fetch, link, redirect);
    }

    @Override
    public boolean isDBExists() {
        return true;
    }

    @Override
    public void clear() throws Exception {
        storeDB.crawlDB.clear();
        storeDB.fetchDB.clear();
        storeDB.linkDB.clear();
        storeDB.redirectDB.clear();
    }

    @Override
    public Generator createGenerator() {
        return new StoreDBGenerator(storeDB);
    }

    @Override
    public void open() throws Exception {
    }

    @Override
    public void close() throws Exception {
    }

    @Override
    public void inject(CrawlDatum datum, boolean force) throws Exception {
        String key = datum.key();
        if (!force) {
            if (storeDB.crawlDB.containsKey(key)) {
                return;
            }
        }
        storeDB.crawlDB.put(key, gson.toJson(datum));
    }

    @Override
    public void inject(CrawlDatums datums, boolean force) throws Exception {
        for (CrawlDatum datum : datums) {
            inject(datum, force);
        }
    }

    @Override
    public void merge() throws Exception {
        LOG.info("start merge");

        /*合并fetch库*/
        LOG.info("merge fetch database");
        for (Map.Entry<String, String> fetchEntry : storeDB.fetchDB.entrySet()) {
            storeDB.crawlDB.put(fetchEntry.getKey(), fetchEntry.getValue());
        }

        /*合并link库*/
        LOG.info("merge link database");
        for (String key : storeDB.linkDB.keySet()) {
            if (!storeDB.crawlDB.containsKey(key)) {
                storeDB.crawlDB.put(key, storeDB.linkDB.get(key));
            }
        }

        LOG.info("end merge");

        storeDB.fetchDB.clear();
        LOG.debug("remove fetch database");
        storeDB.linkDB.clear();
        LOG.debug("remove link database");

    }

    @Override
    public void initSegmentWriter() throws Exception {
    }

    @Override
    public synchronized void writeFetchSegment(CrawlDatum fetchDatum) throws Exception {
        storeDB.fetchDB.put(fetchDatum.key(), gson.toJson(fetchDatum));
    }

    @Override
    public synchronized void writeParseSegment(CrawlDatums parseDatums) throws Exception {
        for (CrawlDatum datum : parseDatums) {
            storeDB.linkDB.put(datum.key(), gson.toJson(datum));
        }
    }

    @Override
    public void closeSegmentWriter() throws Exception {
    }

    @Data
    @Builder
    @AllArgsConstructor
    public static class StoreDB {
        protected Map<String, String> crawlDB = new HashMap<>();
        protected Map<String, String> fetchDB = new HashMap<>();
        protected Map<String, String> linkDB = new HashMap<>();
        protected Map<String, String> redirectDB = new HashMap<>();
    }


    public static class StoreDBGenerator extends Generator {

        private Gson gson = new Gson();

        Iterator<Map.Entry<String, String>> iterator;

        public StoreDBGenerator(StoreDB storeDB) {
            this.iterator = storeDB.crawlDB.entrySet().iterator();
        }

        @Override
        public CrawlDatum nextWithoutFilter() throws Exception {
            if (iterator.hasNext()) {
                CrawlDatum datum = gson.fromJson(iterator.next().getValue(), CrawlDatum.class);
                return datum;
            } else {
                return null;
            }
        }

        @Override
        public void close() throws Exception {

        }
    }
}
